/** \file    optic_flow.cpp
  * \brief   Makes transform more precise by using direct pixels-based motion estimation method.
  * \date    2002, 2003
  * \author  Fedor Chelnokov
  * \version $Header: /home/cscom1r4/cvsroot/cs/imglib/src/optic_flow.cpp,v 1.8 2007/10/07 14:40:05 cscom1r4 Exp $
  */

#include <glib/g2/p2p_matcher.h>
#include <glib/g2/p2l_matcher.h>
#include <glib/g2/translation.h>
#include <glib/g2/rigid.h>
#include <glib/g2/rigid_scale.h>
#include <glib/g2/affine.h>
#include <glib/g4/vector.h>
#include <glib/g4/smatrix.h>
#include <glib/g2/io.h>
#include <glib/g2/optimize.h>

#include "imglib/optic_flow.h"
#include "imglib/imageclass.h"
//#include "gslmath/motionfit.h"

#include "mswin/timing.h"

#include "imglib/type.h"
#include "typelib/affine.h"

using std::clog;
using std::endl;

std::ostream & operator <<(std::ostream & os, const affine & a)
{
//	img::QuadraticForm t;
//	t.assign(a);
  g2::affinef t(g2::matrixf(a.xx, a.xy, a.yx, a.yy), g2::vectorf(a.x, a.y));
	return os << t;
}

namespace img {

template <class T>
class transform_solver {
public:
	typedef typename optic_flow<T>::pixel0 pixel0;

public:
	virtual ~transform_solver() {}
	virtual const char * name() const = 0;
	virtual void alloc(unsigned n) = 0;
	virtual void set_lhs(size_t r0, size_t x, size_t y, 
		const float_color & I, const float_color & Ix, const float_color & Iy) = 0;
	virtual void prepare_lhs(
    const std::vector<pixel0> & frame0_data) = 0;
	virtual double prepare_rhs(
		const std::vector<pixel0> & frame0_data,
		const std::vector<float_color> & frame1_colors) = 0;
	virtual void solve(const mswin::timing_tree * tt = 0) = 0;
	virtual bool uncertain() const = 0;
	virtual void amend(T & t) const = 0;

	//potion of pixels setup
	virtual float alpha() const = 0;
	virtual void alpha(float a) = 0;
};

///implementation independent from matching method
template <class T>
class transform_solver_impl : public transform_solver<T> {
public:
	transform_solver_impl() {
		_alpha = 1.0;
	}
	virtual float alpha() const {
		return _alpha;
	}
	virtual void alpha(float a) {
		_alpha = a;
	}
	virtual void amend(T & t) const {
		t = t * _amendment;
	}
protected:
	///increments measure of correlation between two frames. This measure is
	///1) fast to compute
	///2) alpha independent
	///3) robust to outliers
	void inc_norm(double & chisq, float Ir_squared, float It_squared) {
		chisq += std::min(Ir_squared, It_squared);
	}
	void adjust_alpha(unsigned ignored, unsigned total) {
		clog << "(ignored pixels " << ignored << " / " << total;
		float per = float(ignored) / total;
		int p = clog.precision(3);
		clog << ", " << per*100 << '%';
		clog.precision(p);
		if (per > 0.10) {
			_alpha *= 2;
			clog << ", portion decreased";
		}
		else if (per < 0.02 && _alpha > 1.5) {
			_alpha *= 0.5;
			clog << ", portion increased";
		}
		clog << ") ";
	}
protected:
	float _alpha;
	T _amendment;
};

///////////////////////////
// point-to-line solvers //
///////////////////////////
/*
///point-to-line solver implementation independent from motion model
template <class T>
class p2l_solver_impl0 : public transform_solver_impl<T> {
public:
	virtual void prepare_lhs(const std::vector<pixel0> &) {
		_mfit->prepare();
	}
	virtual double prepare_rhs(
		const std::vector<pixel0> & frame0_data,
		const std::vector<float_color> & frame1_colors)
	{
		unsigned r0 = 0;
		double chisq = 0;
		unsigned ig = 0;
		std::vector<pixel0>::const_iterator p0 = frame0_data.begin();
		std::vector<float_color>::const_iterator i1 = frame1_colors.begin();
		for (; i1 != frame1_colors.end(); ++p0, ++i1, r0+=3) {
			float Ir_squared = p0->squared_Ir;
			float_color It = *i1 - p0->I;
			float It_squared = (float) square(It);
			//robustness
			if (It_squared > _alpha * Ir_squared) {
				_mfit->sety(r0,   0);
				_mfit->sety(r0+1, 0);
				_mfit->sety(r0+2, 0);
				++ig;
			}
			else {
				_mfit->sety(r0,   -It.r);
				_mfit->sety(r0+1, -It.g);
				_mfit->sety(r0+2, -It.b);
			}
			inc_norm(chisq, Ir_squared, It_squared);
		}
		adjust_alpha(ig, frame1_colors.size());
		return chisq;
	}
	virtual bool uncertain() const {
		return transform_err<T>(_amendment, _error).uncertain();
	}
protected:
	void set3cells(size_t r0, size_t c, double v, const float_color & I) {
		_mfit->setX(r0,   c, v*I.r);
		_mfit->setX(r0+1, c, v*I.g);
		_mfit->setX(r0+2, c, v*I.b);
	}
	void set3cells(size_t r0, size_t c, const float_color & I) {
		_mfit->setX(r0,   c, I.r);
		_mfit->setX(r0+1, c, I.g);
		_mfit->setX(r0+2, c, I.b);
	}
	void report() {
		clog << "residual (255-max): " << sqrt(_chisq/_mfit->n()) << "\n";
		clog << "found amendment:\n" << transform_err<T>(_amendment, _error);
	}

protected:
	T _error;
	std::auto_ptr<motionfit> _mfit;
	double _chisq;
};

template <class T, class M>
class p2l_solver_impl : public p2l_solver_impl0<T> {
public:
	virtual const char * name() const {
		return M::name();
	}
	virtual void alloc(unsigned n) {
		_mfit.reset(new motionfit(M::param_count(), 3*n));
	}
	virtual void solve(const mswin::timing_tree * tt = 0) {
		mswin::timing_mark mark(tt, "least-square solving");
		_chisq = _mfit->run();
		_mfit->get_transform(0, M(), _amendment, &_error);
		report();
	}
public:
	float_color Ix(size_t r0) {
		const size_t a0col = M::a0col();
		return float_color(_mfit->X(r0+2, a0col), _mfit->X(r0+1, a0col), _mfit->X(r0, a0col));
	}
	float_color Iy(size_t r0) {
		const size_t b0col = M::b0col();
		return float_color(_mfit->X(r0+2, b0col), _mfit->X(r0+1, b0col), _mfit->X(r0, b0col));
	}
protected:
	void set6cells(size_t r0, size_t c, double v, 
				   const float_color & Ix, const float_color & Iy) {
		assert(M::separable());
		const size_t ycol = M::param_count() / 2;
		set3cells(r0, c, v, Ix);
		set3cells(r0, ycol+c, v, Iy);
	}
	void set6cells(size_t r0, size_t c, 
				   const float_color & Ix, const float_color & Iy) {
		assert(M::separable());
		const size_t ycol = M::param_count() / 2;
		set3cells(r0, c, Ix);
		set3cells(r0, ycol+c, Iy);
	}
};

template <class T>
class translation_solver : public p2l_solver_impl<T, translation_model> {
public:
	virtual void set_lhs(size_t r0, size_t x, size_t y, 
		const float_color & I, const float_color & Ix, const float_color & Iy) 
	{
		set6cells(r0, 0, Ix, Iy);
	}
};

template <class T>
class rigid_zoom_solver : public p2l_solver_impl<T, rigid_zoom_model> {
public:
	virtual void set_lhs(size_t r0, size_t x, size_t y, 
		const float_color & I, const float_color & Ix, const float_color & Iy) 
	{
		set3cells(r0, 0, Ix*float(x) + Iy*float(y));
		set3cells(r0, 1, Ix*float(y) - Iy*float(x));
		set3cells(r0, 2, Ix);
		set3cells(r0, 3, Iy);
	}
};

template <class T>
class rigid_solver : public rigid_zoom_solver<T> {
public:
	virtual const char * name() const {
		return rigid_model::name();
	}
	virtual void solve(const mswin::timing_tree * tt = 0) {
		rigid_zoom_solver<T>::solve(tt);
		_amendment.rigid();
		_error.rigid();
	}
};

template <class T>
class affine_solver : public p2l_solver_impl<T, affine_model> {
public:
	virtual void set_lhs(size_t r0, size_t x, size_t y, 
		const float_color & I, const float_color & Ix, const float_color & Iy) 
	{
		set6cells(r0, 0, x, Ix, Iy);
		set6cells(r0, 1, y, Ix, Iy);
		set6cells(r0, 2,    Ix, Iy);
	}
};

template <class T>
class quadratic_solver : public p2l_solver_impl<T, quadratic_model> {
public:
	virtual void set_lhs(size_t r0, size_t x, size_t y, 
		const float_color & I, const float_color & Ix, const float_color & Iy) 
	{
		set6cells(r0, 0, x, Ix, Iy);
		set6cells(r0, 1, y, Ix, Iy);
		set6cells(r0, 2,    Ix, Iy);
		set6cells(r0, 3, x*x, Ix, Iy);
		set6cells(r0, 4, x*y, Ix, Iy);
		set6cells(r0, 5, y*y, Ix, Iy);
	}
};

template <>
class quadratic_solver<affine> : public affine_solver<affine> {
public:
	quadratic_solver<affine>()
		{ assert(false); }
};
*/
////////////////////////////////
// fast point-to-point solver //
////////////////////////////////

/// 1 / (127.5^2 + 127.5^2)
const float r_max_grad2 = 2.0f / (255 * 255);

//point-to-point solver implementation independent from motion model
class solver_impl : public transform_solver_impl<affine> {
public:
	virtual void alloc(unsigned n) {
		_grad_data.resize(3*n);
	}
	virtual void set_lhs(size_t r0, size_t /*x*/, size_t /*y*/, 
		const float_color & /*I*/, const float_color & Ix, const float_color & Iy)
	{
		set_grad(r0,   float_point(Ix.r, Iy.r));
		set_grad(r0+1, float_point(Ix.g, Iy.g));
		set_grad(r0+2, float_point(Ix.b, Iy.b));
	}
	virtual bool uncertain() const {
		return false;
	}
protected:
	void report() const {
		clog << "found amendment:\n" << _amendment << "\n";
	}

private:
	void set_grad(size_t r, const float_point & Ir) {
		grad & g = _grad_data[r];
		float val2 = (float) square(Ir);
		if (val2 > 0.0001)
    {
      g.weight = val2 * r_max_grad2;
			g.rval = sqrt(1.0f / val2);
      g.dir = Ir * g.rval;
    }
		else
    {
      g.weight = 0;
			g.rval = 0;
      g.dir = float_point(0, 0);
    }
		
	}
protected:
	struct grad {
		float_point dir;
		float rval;
    float weight;
	};
	std::vector<grad> _grad_data;
};

template <class M>
class solver;

template <class T>
class solver< g2::p2p_matcher<T> > : public solver_impl 
{
public:
  typedef T transform;
public:
	virtual const char * name() const
		{ return transform::name(); }
	virtual void prepare_lhs(const std::vector<pixel0> &)
		{ }
	virtual double prepare_rhs(
		const std::vector<pixel0> & frame0_data,
		const std::vector<float_color> & frame1_colors)
	{
		_m.clear();
		unsigned r0 = 0;
		double chisq = 0;
		unsigned ig = 0;
		std::vector<pixel0>::const_iterator p0 = frame0_data.begin();
		std::vector<float_color>::const_iterator i1 = frame1_colors.begin();
		for (; i1 != frame1_colors.end(); ++p0, ++i1, r0+=3) {
			float Ir_squared = p0->squared_Ir;
			float_color It = *i1 - p0->I;
			float It_squared = (float) square(It);
			//robustness
			if (It_squared <= _alpha * Ir_squared) {
				float_point fp(p0->p);
				add_point(r0,   fp, It.r);
				add_point(r0+1, fp, It.g);
				add_point(r0+2, fp, It.b);
			}
      else
				++ig;
      inc_norm(chisq, Ir_squared, It_squared);
		}
		adjust_alpha(ig, frame1_colors.size());
		return chisq;
	}
	virtual void solve(const mswin::timing_tree * = 0) {
    transform t;
		_m.get_transform(t);
    g2::affinef a = t;
    _amendment.x = a.d.x;
    _amendment.y = a.d.y;
    _amendment.xx = a.r.xx;
    _amendment.xy = a.r.xy;
    _amendment.yx = a.r.yx;
    _amendment.yy = a.r.yy;
		report();
	}
private:
	void add_point(unsigned r, const float_point & fp, float It) {
		const grad & g = _grad_data[r];
    if (g.weight > 0)
		  _m.match(fp, fp - g.dir * (It * g.rval), g.dir, g.weight);
	}
private:
	g2::p2p_matcher<T> _m;
};

template <class T>
class solver< g2::p2l_matcher<T> > : public solver_impl 
{
public:
  typedef T transform;
public:
	virtual const char * name() const
		{ return transform::name(); }
	virtual void prepare_lhs(const std::vector<pixel0> & frame0_data)
	{ 
		_m.clear();
		unsigned r0 = 0;
		std::vector<pixel0>::const_iterator p0 = frame0_data.begin();
		for (; p0 != frame0_data.end(); ++p0, r0+=3) {
			float_point fp(p0->p);
			add_point(r0,   fp);
			add_point(r0+1, fp);
			add_point(r0+2, fp);
		}
  }
	virtual double prepare_rhs(
		const std::vector<pixel0> & frame0_data,
		const std::vector<float_color> & frame1_colors)
	{
		_m.clear_p();
		unsigned r0 = 0;
		double chisq = 0;
		unsigned ig = 0;
		std::vector<pixel0>::const_iterator p0 = frame0_data.begin();
		std::vector<float_color>::const_iterator i1 = frame1_colors.begin();
		for (; i1 != frame1_colors.end(); ++p0, ++i1, r0+=3) {
			float Ir_squared = p0->squared_Ir;
			float_color It = *i1 - p0->I;
			float It_squared = (float) square(It);
			if (It_squared > _alpha * Ir_squared)
      {
        It = float_color(0, 0, 0); //robustness
				++ig;
      }
			float_point fp(p0->p);
			change_point(r0,   fp, It.r);
			change_point(r0+1, fp, It.g);
			change_point(r0+2, fp, It.b);
      inc_norm(chisq, Ir_squared, It_squared);
		}
		adjust_alpha(ig, frame1_colors.size());
		return chisq;
	}
	virtual void solve(const mswin::timing_tree * = 0) {
    transform t;
		_m.get_transform(t);
    g2::affinef a = t;
    _amendment.x = a.d.x;
    _amendment.y = a.d.y;
    _amendment.xx = a.r.xx;
    _amendment.xy = a.r.xy;
    _amendment.yx = a.r.yx;
    _amendment.yy = a.r.yy;
		report();
	}
private:
	void add_point(unsigned r, const float_point & fp) {
		const grad & g = _grad_data[r];
    if (g.weight > 0)
		  _m.match(fp, 0, g.dir, g.weight);
	}
	void change_point(unsigned r, const float_point & fp, float It) {
		const grad & g = _grad_data[r];
    if (g.weight > 0)
		  _m.match_p(fp, scalar_product(fp, g.dir) - It * g.rval, g.dir, g.weight);
	}
private:
	g2::p2l_matcher<T> _m;
};

///////////////////////
// optic_flow itself //
///////////////////////

template <class T>
optic_flow<T>::optic_flow()
{
	_extra_solver = _solver = 0;
}

template <class T>
optic_flow<T>::~optic_flow()
{
	delete _solver;
	delete _extra_solver;
}

template <class T>
void optic_flow<T>::transform(const T & t, double discrepancy)
{
	_solution.transform = t;
	_solution.amendment.identity();
	_solution.discrepancy = discrepancy;
}

template <class T>
void optic_flow<T>::new_p2l_solver(TransformClass model)
{
	assert(!_solver);
	switch (model) {
	case TRANSFORM_TRANSLATION:
		_solver = new solver<g2::p2l_matcher<g2::translationf> >;
		break;
	case TRANSFORM_RIGID:
		_solver = new solver<g2::p2l_matcher<g2::rigidf> >;
		break;
	case TRANSFORM_RIGID_ZOOM:
		_solver = new solver<g2::p2l_matcher<g2::rigid_scalef> >;
		break;
	case TRANSFORM_AFFINE:
		_solver = new solver<g2::p2l_matcher<g2::affinef> >;
		break;
//	case TRANSFORM_QUADRATIC:
//		_solver = new quadratic_solver<T>;
//		break;
	default:
		assert(false);
	}
}

template <>
void optic_flow<affine>::new_p2p_solver(TransformClass model)
{
	assert(!_extra_solver);
	switch (model) {
	case TRANSFORM_TRANSLATION:
    _extra_solver = new solver<g2::p2p_matcher<g2::translationf> >;
		break;
	case TRANSFORM_RIGID:
		_extra_solver = new solver<g2::p2p_matcher<g2::rigidf> >;
		break;
	case TRANSFORM_RIGID_ZOOM:
		_extra_solver = new solver<g2::p2p_matcher<g2::rigid_scalef> >;
		break;
	case TRANSFORM_AFFINE:
		_extra_solver = new solver<g2::p2p_matcher<g2::affinef> >;
		break;
	default:
		assert(false);
	}
}

//template <>
//void optic_flow<QuadraticForm>::new_p2p_solver(TransformClass model)
//{
//	assert(false);
//}

template <class T>
bool optic_flow<T>::prepare(const T & t, 
	const float_image & frame0, const bitmap & mask0, TransformClass model, bool point2point)
{
	transform(t);

	_w = frame0.width();
	_h = frame0.height();
	//do not trust 5-pixel frame bound
	rect<float> smaller_frame(5, 5, float(_w)-6, float(_h)-6);
	if (smaller_frame.empty())
		return false;

	bitmap::const_iterator b;
	_frame0_data.clear();
	_frame0_data.reserve(mask0.set_bit_count());

	for (b = mask0.begin(); b != mask0.end(); ++b) {
		if (!*b)
			continue;
		const size_t x = b.x(), y = b.y();
		point<float> tp = _solution.transform(float(x), float(y));
		if (x!=0 && x!=_w-1 && y!=0 && y!=_h-1 && smaller_frame.inside(tp))
			_frame0_data.push_back(pixel0(frame0(x,y), x, y));
	}

	float alpha = 1.0;
	if (_solver) {
		alpha = _solver->alpha();
		clog << "initial portion coefficient: " << alpha << '\n';
		delete _solver;
		_solver = 0;
	}
	delete _extra_solver;
	_extra_solver = 0;
	if (_frame0_data.empty())
		return false;

	new_p2l_solver(model);
	_solver->alpha(alpha);
	if (point2point) {
		new_p2p_solver(model);
		_extra_solver->alpha(alpha);
	}

	_solver->alloc(mask_size());
	if (_extra_solver)
		_extra_solver->alloc(mask_size());

	size_t r = 0;
	float_color Ix, Iy;
	std::vector<pixel0>::iterator p0 = _frame0_data.begin();
	for (; p0 != _frame0_data.end(); ++p0) {
		const size_t x = p0->p.x, y = p0->p.y;
		Ix = 0.5 * (frame0(x+1,y) - frame0(x-1,y));
		Iy = 0.5 * (frame0(x,y+1) - frame0(x,y-1));
		p0->squared_Ir = float(square(Ix) + square(Iy));

		_solver->set_lhs(r, x, y, p0->I, Ix, Iy);
		if (_extra_solver)
			_extra_solver->set_lhs(r, x, y, p0->I, Ix, Iy);
		r += 3;
	}
	_solver->prepare_lhs(_frame0_data);
	if (_extra_solver)
		_extra_solver->prepare_lhs(_frame0_data);
	return true;
}

template <class T>
double optic_flow<T>::discrepancy(const float_image & frame1, const mswin::timing_tree * tt)
{
    assert(_solver);
    if (!_solver)
        return 0;

	mswin::timing_mark mark(tt, "discrepancy calculating");

	T t = _solution.transform;
	t.translate(0.5, 0.5); //(+0.5f, +0.5f) is required by bilinear interpolater
	rect<float> frame_rect(0.5f, 0.5f, float(_w)-0.5f, float(_h)-0.5f);
	bitmap::iterator b;

	unsigned n = 0;
	std::vector<float_color> frame1_colors;
	frame1_colors.reserve(mask_size());
	{
		mswin::timing_mark mark(tt, "bilinear diff.");
		std::vector<pixel0>::const_iterator p0 = _frame0_data.begin();
		for (; p0 != _frame0_data.end(); ++p0) {
			float_color color1;
			if (bilinear(color1, frame1, t(p0->p))) {
				frame1_colors.push_back(color1);
				n += 3;
			} else
				frame1_colors.push_back(float_color()); //it will stabilize solution
		}
	}

	{
		mswin::timing_mark mark(tt, "setting rhs");
		return _solution.discrepancy = 
			sqrt(_solver->prepare_rhs(_frame0_data, frame1_colors) / n);
	}
}

template <class T>
void optic_flow<T>::optimize(const float_image & frame1, unsigned iterations, 
	const mswin::timing_tree * tt)
{
	//_control87(_EM_INVALID | _EM_INEXACT | _EM_UNDERFLOW, _MCW_EM);
	if (!_solver)
		return;

	clog << "\nMaking more precise transformation on pixel level ("
		 << _solver->name() << " model)...\n";

	{
		mswin::timing_mark mark(tt, "optic flow optimizing");
		optimize_single_solver(frame1, iterations, tt);
	}

	if (_extra_solver) {
		clog << "Switching over extra optic flow solver\n";
		mswin::timing_mark mark(tt, "optic flow optimizing (extra solver)");

		std::swap(_solver, _extra_solver);
		_solution.discrepancy = -1;
		optimize_single_solver(frame1, iterations, tt);
		std::swap(_solver, _extra_solver);
	}
}

template <class T>
void optic_flow<T>::optimize_single_solver(const float_image & frame1, unsigned iterations, 
	const mswin::timing_tree * tt)
{
	assert(_solver);

	if (!_solution.discrepancy_found())
		clog << "Initial discrepancy is " << discrepancy(frame1, tt) << "\n";

	float initial_discrepancy = (float) _solution.discrepancy;
	solution best_solution = _solution;
	unsigned successive_bad_iterations = 0;

  unsigned i;
	for (i = 0; i < iterations; ++i) 
	{
		clog << "iteration #" << i << ", ";

		_solver->solve(tt);
		_solver->amend(_solution.transform);
		_solver->amend(_solution.amendment);

		clog << "corrected transformation with discrepancy " 
			 << discrepancy(frame1, tt) << " is\n" << _solution.transform << endl;

		if (best_solution.replace_if_better(_solution)) {
			successive_bad_iterations = 0;
			/*for (unsigned j = 0; j < iterations; ++j) {
				_solver->amend(_solution.transform);
				_solver->amend(_solution.amendment);
				clog << "discrepancy after same amendment: " << discrepancy(frame1, tt) << endl;
				if (!best_solution.replace_if_better(_solution)) {
					_solution = best_solution;
					break;
				}
			}*/
		}
		else
			if (++successive_bad_iterations == 5 || _solver->uncertain()) {
				++i;
				break;
			}
	}

	_solution.replace_if_better(best_solution);
	clog << "best transformation (from iteration " << int(i-successive_bad_iterations-1) << ") with discrepancy " 
		 <<	_solution.discrepancy << " is\n" << _solution.transform << "\n"
			"discrepancy lowering gain: " << initial_discrepancy - _solution.discrepancy;
	
	int p = clog.precision(3);
	clog << " (" << 100 * (initial_discrepancy - _solution.discrepancy) / initial_discrepancy << "%)\n"
			"overall amendment:\n" << _solution.amendment << endl;
	clog.precision(p);
}

//template class optic_flow<QuadraticForm>;
template class optic_flow<affine>;

} //namespace img
